# -*- coding: utf-8 -*-
# created on 2016/3/14
from sympy.core.add import _unevaluated_Add
from sympy.abc import x
from sympy import S, Add, degree, And, Mul, Pow, sympify, Rational, Abs, sqrt, Interval, Union, FiniteSet, Set, Symbol
from sympy.functions.elementary.trigonometric import TrigonometricFunction


def get_all_child(expr, func):
    childrens = []
    if func(expr):
        return [expr]

    if hasattr(expr, "args"):
        for arg in expr.args:
            if func(arg):
                # if arg.is_Pow and arg.args[1] != -1:
                childrens.append(arg)
            else:
                childrens.extend(get_all_child(arg, func))

        return childrens
    else:
        return childrens


def default_symbol(expr, args=None, func=None):
    if args is None:
        args = []
    if len(args) > 1:
        return args[1].sympify()
    elif func and func.target:
        return func.target[0].sympify()
    elif x in expr.free_symbols:
        return x
    else:
        try:
            return expr.free_symbols.pop()
        except KeyError:
            # 如果 expr 是常数
            return x


def target_symbol(poly, args, func):
    if len(args) > 1:
        target = args[-1].sympify()
    elif func.target:
        target = func.target[0].sympify()
    elif x in poly.free_symbols:
        target = [sym for sym in poly.free_symbols if sym != x]
    else:
        target = poly.free_symbols.pop()
    return target


def free_symbols(item):
    return list(set([sym for expr in item for sym in expr.free_symbols]))


def peifang(expr, sym):
    expr = expr.expand()
    if expr.is_Pow and expr.args[1] == 2:
        return expr

    if expr.is_Add:
        a, h = expr.as_independent(sym)
    else:
        a, h = S.Zero, expr

    if safe_degree(h, sym) == 2 and h.is_Add:
        coeff_yici = 0
        coeff_erci = 0
        for arg in h.args:
            if safe_degree(arg, sym) == 1:
                coeff_yici = arg / sym
            elif safe_degree(arg, sym) == 2:
                coeff_erci = arg / (sym ** 2)

        v = coeff_yici / (2 * coeff_erci)
        return coeff_erci * (sym + coeff_yici / (2 * coeff_erci)) ** 2 + a - coeff_erci * v ** 2

    return None


def out_sqrt(sqrt_expr):
    assert sqrt_expr.is_Pow and sqrt_expr.args[1] == Rational(1, 2)
    expr = peifang(sqrt_expr.args[0], sqrt_expr.args[0].free_symbols.pop())
    if expr is None:
        return sqrt_expr
    result_expr = 1
    if expr.is_Mul:
        for item in expr.args:
            if item.is_Pow:
                result_expr = result_expr * Abs(item.args[0])
            else:
                result_expr = result_expr * sqrt(item)
    elif expr.is_Pow:
        result_expr = result_expr * Abs(expr.args[0])
    elif isinstance(expr, Symbol):
        result_expr = result_expr * Abs(expr.args[0])
    else:
        raise ValueError("unexpected situation")
    return result_expr


class ReplaceAll(object):
    def __init__(self, symbol="t"):
        self.num = 0
        self.symbol = symbol
        self.symbol_value = {}

    def replace(self, expr, func):
        if func(expr):
            self.num += 1
            symbol = sympify("%s%d" % (self.symbol, self.num))
            self.symbol_value[symbol] = expr
            return symbol

        if expr.is_Add:
            args = [self.replace(arg, func) for arg in expr.args]
            return Add._from_args(args)
        elif expr.is_Mul:
            args = [self.replace(arg, func) for arg in expr.args]
            return Mul._from_args(args)
        elif expr.is_Pow:
            return Pow(self.replace(expr.args[0], func), self.replace(expr.args[1], func))
        elif expr.is_Function:
            return expr.func(*[self.replace(arg, func) for arg in expr.args])
        else:
            return expr


def safe_degree(expr, symbol, nan_value=None):
    try:
        return degree(expr, symbol)
    except Exception:
        return nan_value


def simplify_ineq_solveset(ineq_solveset):
    if isinstance(ineq_solveset, And):
        newset = [arg for arg in ineq_solveset.args if all([term.is_finite or term.is_Symbol for term in arg.args])]
        return And._from_args(newset)
    else:
        return ineq_solveset


def get_pow(expr, sym):
    if not expr.has(sym):
        return 0

    a, h = expr.as_independent(sym)
    if h.is_Pow and h.args[0].has(sym):
        return h.args[1]
    elif h.is_Add:
        return max(*[get_pow(arg, sym) for arg in h.args])
    elif h.is_Mul:
        return sum([get_pow(arg, sym) for arg in h.args])
    elif h.is_Symbol:
        return 1
    else:
        return -1


def interval_2_ineqs(interval, e):
    """将区间转换成 Mathematica 的表达式 e.g. '-3 <= x < 0 || x > 1'"""
    if isinstance(interval, Interval):
        if interval.left != -S.Infinity and interval.right != S.Infinity:
            res = [[interval.left, '<' if interval.left_open else '<=', e],
                   [e, '<' if interval.right_open else '<=', interval.right]]
        elif interval.left == -S.Infinity:
            res = [[e, '<' if interval.right_open else '<=', interval.right]]
        else:
            res = [[e, '>' if interval.left_open else '>=', interval.left]]

    elif isinstance(interval, Union):
        res = [ineq for item in interval.args for ineq in interval_2_ineqs(item, e)]
    else:
        raise ValueError('interval must be of type Interval or Union')

    return res


def independent(expr, symbol):
    if expr.is_zero:
        return S.Zero, S.Zero

    # expr = expr.expand()
    if expr.is_Add():
        depend = indep = []
        for arg in expr.args:
            if arg.has(symbol):
                depend.append(arg)
            else:
                indep.append(arg)

        return Add(*indep), _unevaluated_Add(*depend)
    else:
        return S.Zero, expr


def simple_invert(f1, f2):
    if hasattr(f1, 'inverse') and not isinstance(f1, TrigonometricFunction) and f1.is_Function:
        return f1.args[0], f1.inverse()(f2)


def to_set(item):
    return item if isinstance(item, Set) else FiniteSet(item)


def merge_known(known_value, symbol_value):
    known_key, known_value = list(known_value.items())[0]
    new_key, new_value = list(symbol_value.items())[0]

    # to_set = lambda x: x if isinstance(x, Set) else FiniteSet(x)
    if isinstance(known_value, (Interval, Union)):
        known_value_pairs = [{known_key: known_value}]
    elif isinstance(known_value, (FiniteSet, list, tuple)):
        known_value_pairs = [
            dict([(k, to_set(v)) for k, v in zip(known_key, value)]) if isinstance(known_key, tuple) else {
                known_key: to_set(value)} for value in known_value]
    else:
        known_value_pairs = [{known_key: FiniteSet(known_value)}]

    if isinstance(new_value, (Interval, Union)):
        new_value_pairs = [{new_key: new_value}]
    elif isinstance(known_value, (FiniteSet, list, tuple)):
        new_value_pairs = [dict([(k, to_set(v)) for k, v in zip(new_key, value)]) if isinstance(new_key, tuple) else {
            new_key: to_set(value)} for value in new_value]
    else:
        new_value_pairs = [{new_key: FiniteSet(new_value)}]

    all_values = []
    for p1 in known_value_pairs:
        for p2 in new_value_pairs:
            all_value = dict(p1)
            for k, v in p2.items():
                if k in all_value:
                    if isinstance(p2[k], FiniteSet) and not isinstance(all_value[k], FiniteSet):
                        value = p2[k]
                    else:
                        value = all_value[k].intersect(p2[k])
                else:
                    value = p2[k]

                if isinstance(value, FiniteSet):
                    value = value.args[0]

                all_value[k] = value

            all_value = dict([(k, v.args[0]) if isinstance(v, FiniteSet) else (k, v) for k, v in all_value.items()])
            all_values.append(all_value)

    all_key = tuple(all_values[0].keys())
    all_key = all_key if len(all_key) > 1 else all_key[0]
    all_value = [tuple(list(value.values())) if len(value) > 1 else list(value.values())[0] for value in all_values]
    all_value = all_value[0] if len(all_value) == 1 and isinstance(all_value[0], Set) else FiniteSet(*all_value)
    return {all_key: all_value}


if __name__ == '__main__':
    pass
